import os
import re
import signal
from tempfile import NamedTemporaryFile, mkstemp
from threading import Thread, Event, Lock
from subprocess import Popen, PIPE
# TODO: look into the pty module and openpty(), to address some problems we've seen with popen and subprocess
from getpass import getpass
import paramiko
from pytoss._async import AsyncIterable, AsyncIterableGroup
from pytoss.support import PickledExecution, ASYNCTRACE

def _bytes_from_stream(stream):
    """yields the bytes of a stream one at a time"""
    next_char = stream.read(1)
    while next_char:
        yield next_char
        next_char = stream.read(1)

def create_connection(username, hostname, port = 22, getpass = None, su_source_username = None, su_source_getpass = None):

        if hostname == "localhost":
            # we have a local scenario
        
            if username == os.getlogin():
                # local user, on local machine
                return LocalUserConnection()
            
            else:
                # other user, on local machine
                raise UnsupportedConnectionError("currently cannot support local users other than the current user (use os.getlogin() to get the local username)")
    
        else:
            # we have a remote scenario
            if su_source_username:
                # user on remote machine who needs to be SSHed in
                # and then su'd to before it can do anything
                su_from_connection = RemoteConnection(su_source_username, hostname, getpass = su_source_getpass, port = port)
                return RemoteSuConnection(
                    su_from_connection = su_from_connection,
                    username = username,
                    getpass = getpass,
                )
        
            else:
                # this is a user who can remotely log in to the machine
                # and execute commmands
                return RemoteConnection(
                    username = username,
                    hostname = hostname,
                    getpass = getpass,
                    port = port,
                )


class SFTPFileAdapter(object):
    
    def __init__(self, sftp_file, filepath):
        self.__sftp_file = sftp_file
        self.__filepath = filepath
    
    def __getattr__(self, attr_name):
        if attr_name == "name":
            return self.__filepath
        return getattr(self.__sftp_file, attr_name)

class GeneratorStream(object):
    
    def __init__(self, generator, input_handler, kill_handler):
        self.__async_iterable = AsyncIterable( generator )
        self.__input_handler = input_handler
        self.__kill_handler = kill_handler
        self.__input_lock = Lock()
    
    def __iter__(self):
        for obj in self.__async_iterable:
            yield obj
    
    def inject(self, obj):
        # TODO: this may need some kind of special object (None may not be appropriate) to indicate when no more input will be provided?
        self.__input_lock.acquire()
        try:
            self.__input_handler( obj )
        finally:
            self.__input_lock.release()
        return self
    
    def drain(self):
        for dontcare in self:
            pass
    
    def kill(self):
        self.__kill_handler()
        self.__async_iterable.kill()


class CommandStream(object):
    
    def __init__(self, stdout_generator, stderr_generator, input_handler, kill_handler):
        self.__async_iterable = AsyncIterableGroup( stdout_generator, stderr_generator )
        self.__input_handler = input_handler
        self.__kill_handler = kill_handler
        self.__input_lock = Lock()
    
    # FIXME: make a "bytes" iterator
    # FIXME: yield subclasses of tuple with field names
    
    def __iter__(self):
        """yields each byte of the stream, and also
        (as a convenience) each line of the stream"""
        pending_stdout_line = ""
        pending_stderr_line = ""
        for stdout_byte, stderr_byte in self.__async_iterable:
            stdout_line = None
            stderr_line = None
            if stdout_byte:
                if stdout_byte == "\n":
                    stdout_line = pending_stdout_line
                    pending_stdout_line = ""
                else:
                    pending_stdout_line += stdout_byte
            if stderr_byte:
                if stderr_byte == "\n":
                    stderr_line = pending_stderr_line
                    pending_stderr_line = ""
                else:
                    pending_stderr_line += stderr_byte
            
            # FIXME: rather than an N-tuple, use the 
            yield stdout_line, stderr_line, stdout_byte, stderr_byte
    
    def inject(self, buffer_str):
        self.__input_lock.acquire()
        try:
            self.__input_handler( buffer_str )
        finally:
            self.__input_lock.release()
        return self
    
    def drain(self):
        for dontcare in self:
            pass
    
    def kill(self):
        self.__kill_handler()
        self.__async_iterable.kill()


class UnsupportedConnectionError(Exception): pass
class AuthenticationError(Exception): pass


class NonzeroReturnStatus(Exception):
    
    code = property(lambda self: self.__code)
    command = property(lambda self: self.__command)
    
    def __init__(self, code, command, message = None):
        if message and len(message.strip()) > 0:
            Exception.__init__(self, "'%s' returned exit status %s (%s)" % (command, code, message))
        else:
            Exception.__init__(self, "'%s' returned exit status %s" % (command, code))
        self.__code = code
        self.__command = command
    
    def __str__(self):
        if self.message and len(self.message) > 0:
            return "'%s' returned exit status %s (%s)" % (self.command, self.code, self.message)
        else:
            return "'%s' returned exit status %s" % (self.command, self.code)


class AbstractConnection(object):
    
    def create_command_stream(self, command):
        """return an iterable that sequences over (stdout, stderr)"""
        raise NotImplementedError
        
    def create_generator_stream(self, generator, virtualenv):
        """return an async iterable that sequences over (obj) and raises
        any remote exceptions as PickledExecutionExceptions"""
        raise NotImplementedError
    
    def open_file(self, filepath, mode = "w+"):
        """opens a stream to the given file on the connection"""
        raise NotImplementedError
    
    def remove_file(self, filepath):
        """removes a file on the given connection"""
        raise NotImplementedError
    
    def open_tempfile(self):
        """opens a temp file on the given connection"""
        raise NotImplementedError
    
    def make_directory(self, dirpath):
        """creates a directory on the given connection"""
        raise NotImplementedError
    
    def remove_directory(self, dirpath):
        """removes a directory on the given connection"""
        raise NotImplementedError

class LocalUserConnection(AbstractConnection):
    
    def __init__(self):
        AbstractConnection.__init__(self)
    
    def open_file(self, filepath, mode = "w+"):
        return open(filepath, mode)
    
    def open_tempfile(self):
        filename = mkstemp()[1]
        return self.open_file(filename)
    
    def create_command_stream(self, command):

        pipe = Popen([command], shell = True, close_fds = True, stdout = PIPE, stderr = PIPE, stdin = PIPE)
        done_event = Event()
        
        def watch_stdout():
            
            # yield each line from the stream
            for byte in _bytes_from_stream(pipe.stdout):
                yield byte
                
            # raise exceptions on bad return codes
            done_event.wait()
            code = pipe.returncode
            if code != 0:
                raise NonzeroReturnStatus(code, command)

        def watch_stderr():
            
            # yield each line from the stream
            for byte in _bytes_from_stream(pipe.stderr):
                yield byte
                
            # raise exceptions on bad return codes
            done_event.wait()
            code = pipe.returncode
            if code != 0:
                raise NonzeroReturnStatus(code, command)
        
        # hack to reliably watch for the process to finish
        # (for some reason, pipe.wait() could sometimes fail with no child process)
        def watch_status():
            try:
                if pipe.returncode is None:
                    pipe.wait()
            except OSError, e:
                # TODO: should this do anything? the process is done at this point either way
                pass
            done_event.set()
        
        Thread(target = watch_status).start()
        
        # set up an input handler
        def input_handler(buffer_str):
            pipe.stdin.write( buffer_str )
            
        def kill_handler():
            if pipe.returncode is None:
                os.kill( pipe.pid, signal.SIGKILL )

        # FIXME: somehow needs to register an 'abort' method that cuts off all streams
        return CommandStream(watch_stdout, watch_stderr, input_handler, kill_handler)
        
    def create_generator_stream(self, generator, virtualenv = None):
        
        # create the pickled execution for this generator
        pickled_execution = PickledExecution(generator)
        
        # write the code that must be executed
        temp_code_file = NamedTemporaryFile()
        temp_code_file.write( pickled_execution.code )
        temp_code_file.flush()
        
        # execute it via the shell
        if virtualenv == None:
            python_cmd = "python %s" % temp_code_file.name
        else:
            python_cmd = "%s/bin/python %s" % (virtualenv, temp_code_file.name)
        pipe = Popen([python_cmd], shell = True, close_fds = True, stdout = PIPE, stderr = PIPE, stdin = PIPE)
        
        # asynchronously make sure the temp file is deleted after finishing the command
        def close_tempfile_on_completion():
            pipe.wait()
            temp_code_file.close()
            
        Thread(target = close_tempfile_on_completion).start()
        
        # set up an input handler
        def input_handler(obj):
            pickled_execution.write_object_to_stdin(obj, pipe.stdin)
            
        def kill_handler():
            os.kill( pipe.pid, signal.SIGKILL )
        
        # do the iteration over the resulting python execution
        return GeneratorStream( lambda: pickled_execution.get_response_iterator(pipe.stdout, pipe.stderr), input_handler, kill_handler )


class RemoteConnection(AbstractConnection):
    
    def __init__(self, username, hostname, getpass, port):
        AbstractConnection.__init__(self)
        
        # keep track of our connection metadata
        self.__username = username
        self.__hostname = hostname
        self.__port = int(port)
        self.__getpass_callback = getpass
        
        # create our internal ssh client
        self.__ssh_client = paramiko.SSHClient()
        self.__ssh_client.load_system_host_keys()
        self.__ssh_client.set_missing_host_key_policy(paramiko.WarningPolicy)
        self.__sftp_client = None
    
    def __del__(self):
        self.__ssh_client.close()
        if self.__sftp_client:
            self.__sftp_client.close()
    
    def _ensure_sftp_is_ready(self):
        if not self.__sftp_client:
            self._ensure_connected_and_authenticated()
            try:
                self.__sftp_client = self.__ssh_client.open_sftp()
            except Exception, e:
                self.__sftp_client.close()
                raise
    
    def _ensure_connected_and_authenticated(self):
        
        # keep track of the last exception we got from paramiko during authentication attempts
        last_ssh_exception = None
        
        # only do this if we are not authenticated yet
        def is_connected_and_authenticated():
            return self.__ssh_client.get_transport() and self.__ssh_client.get_transport().is_authenticated()
        
        # ATTEMPT 1: use local PKeys
        if not is_connected_and_authenticated():
            
            try:
                self.__ssh_client.connect(
                    hostname = self.__hostname,
                    username = self.__username,
                    port = self.__port,
                )
            except paramiko.SSHException, e:
                last_ssh_exception = e
                self.__ssh_client.close()
            
        # ATTEMPT 2: use ssh agent on the local machine
        if not is_connected_and_authenticated():
        
            try:
                agent = paramiko.Agent()
                
                for pkey in agent.get_keys():
                    try:
                        self.__ssh_client.connect(
                            hostname = self.__hostname,
                            username = self.__username,
                            port = self.__port,
                            pkey = pkey,
                            allow_agent = True,
                            look_for_keys = True
                        )
                    except paramiko.SSHException, e:
                        last_ssh_exception = e
                        self.__ssh_client.close()
                        agent.close()
                        
            except paramiko.SSHException, e:
                last_ssh_exception = e
                self.__ssh_client.close()
                agent.close()
            
        # ATTEMPT 3: final attempt, use password
        if not is_connected_and_authenticated():
            
            try:
                if self.__getpass_callback:
                    password = self.__getpass_callback()
                    self.__ssh_client.connect(
                        hostname = self.__hostname,
                        username = self.__username,
                        port = self.__port,
                        password = password
                    )
            except paramiko.SSHException, e:
                last_ssh_exception = e
                self.__ssh_client.close()
        
        # final check
        if not is_connected_and_authenticated() and last_ssh_exception:
            # re-raise if we *still* aren't authenticated
            raise AuthenticationError("failed to authenticate for '%(username)s@%(hostname)s' after trying accessible ssh keys, sshagent, and password callback. (reason: %(paramiko_reason)s)" % dict(
                username = self.__username,
                hostname = self.__hostname,
                paramiko_reason = last_ssh_exception.message,
            ))
    
    def open_file(self, filepath, mode = "w+"):
        self._ensure_sftp_is_ready()
        sftp_file = self.__sftp_client.file( filepath, mode )
        return SFTPFileAdapter(sftp_file, filepath)
    
    def open_tempfile(self):
        temp_filename = None
        for stdout_line, stderr_line, stdout_byte, stderr_byte in self.create_command_stream("python -c 'import tempfile; print tempfile.mkstemp()[1]'"):
            if stdout_line:
                temp_filename = stdout_line
        
        sftp_file = self.open_file(temp_filename, "w+")
        return SFTPFileAdapter(sftp_file, temp_filename)
    
    def create_command_stream(self, command):

        self._ensure_connected_and_authenticated()
        
        # open an SSH shell
        channel = self.__ssh_client.get_transport().open_session()
        
        # get a pseudo-terminal so that certain commands (e.g. scp, rsync, su) behave more nicely
        channel.get_pty()
        
        # grab the streams on that channel
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)

        # create a callback that will execute the command and return its status
        channel.exec_command(command)
        
        stdout_done_event = Event()
        stderr_done_event = Event()
        
        def watch_stdout():
            
            try:
                # yield each line from the stream
                for byte in _bytes_from_stream(stdout):
                    yield byte
                
                # raise exceptions on bad return codes
                ASYNCTRACE("stdout: receiving exit status")
                code = channel.recv_exit_status()
                if code != 0:
                    ASYNCTRACE("stdout: nonzero return code found, raising NonzeroReturnStatus")
                    raise NonzeroReturnStatus(code, command)
                ASYNCTRACE("stdout: no problems (status = 0)")

            finally:
            
                stdout_done_event.set()

        def watch_stderr():
            
            try:
                # yield each line from the stream
                for byte in _bytes_from_stream(stderr):
                    yield byte
                
                # raise exceptions on bad return codes
                ASYNCTRACE("stderr: receiving exit status")
                code = channel.recv_exit_status()
                if code != 0:
                    ASYNCTRACE("stderr: nonzero return code found, raising NonzeroReturnStatus")
                    raise NonzeroReturnStatus(code, command)
                ASYNCTRACE("stderr: no problems (status = 0)")
            
            finally:
                stderr_done_event.set()
                
        def close_channel_when_stdout_and_stderr_are_done():
            ASYNCTRACE("waiting to close channel")
            stdout_done_event.wait()
            stderr_done_event.wait()
            ASYNCTRACE("closing channel")
            channel.close()
            ASYNCTRACE("closed channel")
        
        # be sure to close down the channel when done
        Thread(name = "Thread-close_channel_when_stdout_and_stderr_are_done", target = close_channel_when_stdout_and_stderr_are_done).start()
        
        # set up an input handler
        def input_handler(buffer_str):
            stdin.write( buffer_str )
            
        def kill_handler():
            channel.close()
                
        # FIXME: should also put in stdin, once the IOManager code is in place.  This way you can track what input has gone to the process
        return CommandStream(watch_stdout, watch_stderr, input_handler, kill_handler)

        
    def create_generator_stream(self, generator, virtualenv = None):
        if virtualenv == None:
            python_cmd_format = "python %s"
        else:
            python_cmd_format = "%s/bin/python %s" % virtualenv
        return self._create_generator_stream_with_python_cmd(generator, python_cmd_format)
    
    def _create_generator_stream_with_python_cmd(self, generator, python_cmd_format):
        self._ensure_connected_and_authenticated()
        
        # create the pickled execution for this generator
        pickled_execution = PickledExecution(generator)
        
        # create a temporary file to write to
        temp_code_file = self.open_tempfile()
        temp_code_file.write( pickled_execution.code )
        temp_code_file.close()
        
        # execute it via the shell
        python_cmd = python_cmd_format % temp_code_file.name
        channel = self.__ssh_client.get_transport().open_session()
        channel.get_pty()
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)
        channel.exec_command(python_cmd)
        
        # set up an input handler
        def input_handler(obj):
            pickled_execution.write_object_to_stdin(obj, stdin)
            
        def kill_handler():
            channel.close()
        
        # do the iteration over the resulting python execution
        return GeneratorStream( lambda: pickled_execution.get_response_iterator(stdout, stderr), input_handler, kill_handler )


class RemoteSuConnection(AbstractConnection):
    """a very hackish AbstractConnection that allows a user to be su'ed to"""
    
    def __init__(self, su_from_connection, username, getpass = None):
        self.__su_from_connection = su_from_connection
        self.__username = username
        self.__getpass = getpass
    
    def open_file(self, filepath, mode = "w+"):
        raise NotImplementedError("'su' functionality is experimental and cannot open files as the current user yet")
    
    def open_tempfile(self, filepath, mode = "w+"):
        raise NotImplementedError("'su' functionality is experimental and cannot open files as the current user yet")
    
    def create_command_stream(self, command):
        if "\"" in command:
            raise ValueError("RemoteSuConnection does not yet support the command '%s'" % command)
        
        # wrap the command and execute it on the su_from RemoteConnection
        su_wrapped_command = "su %s -c \"%s\"" % (self.__username, command)
        command_stream = self.__su_from_connection.create_command_stream(su_wrapped_command)
        
        # wrap the CommandStream we got, by watching stdout/stderr
        # for the first password prompt and responding
        password_prompt_pattern = re.compile(".*assword:.*")
        def password_enabled_stream_iteration():
            pending_stdout_line = ""
            pending_stderr_line = ""
            did_respond_to_password = False
            for stdout_line, stderr_line, stdout_byte, stderr_byte in command_stream:
                
                # track each line as it is built
                if stdout_byte:
                    if stdout_byte == "\n":
                        pending_stdout_line = ""
                    else:
                        pending_stdout_line += stdout_byte
                if stderr_byte:
                    if stderr_byte == "\n":
                        pending_stderr_line = ""
                    else:
                        pending_stderr_line += stderr_byte
                
                # respond to passwords
                if not did_respond_to_password:
                    if password_prompt_pattern.search(pending_stdout_line):
                        password = self.__getpass()
                        command_stream.inject(password + "\r\n")
                        did_respond_to_password = True
                    elif password_prompt_pattern.search(pending_stderr_line):
                        password = self.__getpass()
                        command_stream.inject(password + "\r\n")
                        did_respond_to_password = True
                
                # yield the exact set of data that we *would* have yielded
                yield stdout_line, stderr_line, stdout_byte, stderr_byte
        
        # declare a helper class that will wrap this stream
        class _SuCommandStream(object):
            def __init__(self, command_stream, password_stream_iterator):
                self.__command_stream = command_stream
                self.__password_stream_iterator = password_stream_iterator
            def __iter__(self):
                for stdout_line, stderr_line, stdout_byte, stderr_byte in self.__password_stream_iterator():
                    yield stdout_line, stderr_line, stdout_byte, stderr_byte
            def inject(self, *args, **kwargs):
                self.__command_stream.inject(*args,**kwargs)
            def drain(self, *args, **kwargs):
                self.__command_stream.drain(*args,**kwargs)
            def kill(self, *args, **kwargs):
                self.__command_stream.kill(*args,**kwargs)
                
        # return our wrapped CommandStream
        return _SuCommandStream( command_stream, password_enabled_stream_iteration )
    
    def create_generator_stream(self, generator, virtualenv = None):
        raise NotImplementedError("'su' functionality is experimental and cannot execute remote Python code yet")
    